import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import HeteroData
from torch_geometric.nn import GATConv, to_hetero
from torch_geometric.loader import NeighborLoader
from torch_geometric import transforms as T
import numpy as np
import random
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import os

from utils import set_seed, read_yaml, construct_graph_data_hetero

# 设置随机种子确保可复现性
set_seed(0)


class ClickDataGenerator:
    def __init__(self):
        self.user_features = hereto_data["user_features"]
        self.page_features = hereto_data["page_features"]
        self.edge_index = hereto_data["edge_index"]
        self.y = hereto_data["labels"]
        self.train_mask = hereto_data["train_mask"]
        self.valid_mask = hereto_data["valid_mask"]
        self.test_mask = hereto_data["test_mask"]

    def build_hetero_data(self):
        """构建异构图数据"""
        data = HeteroData()
        data["user"].x = self.user_features
        data["page"].x = self.page_features
        data["user", "clicks", "page"].edge_index = self.edge_index
        data["user"].y = self.y
        data["user"].train_mask = self.train_mask
        data["user"].val_mask = self.valid_mask
        data["user"].test_mask = self.test_mask

        # 添加反向边类型
        data = T.ToUndirected()(data)
        return data


# ===================== 异构图神经网络模型 =====================
class GNNEncoder(torch.nn.Module):
    """图神经网络编码器"""

    def __init__(self, hidden_channels, heads=4):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, heads=heads, add_self_loops=False)
        self.conv2 = GATConv((-1, -1), hidden_channels, heads=1, add_self_loops=False)

        self.activation = nn.ELU()
        self.dropout = nn.Dropout(0.2)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = self.activation(x)
        x = self.dropout(x)

        x = self.conv2(x, edge_index)
        return x


class HeteroGNN(torch.nn.Module):
    """异构图神经网络分类器"""

    def __init__(self, metadata, hidden_channels, out_channels, heads=4):
        super().__init__()
        # 初始化编码器
        self.encoder = GNNEncoder(hidden_channels, heads)

        # 转换为异构图卷积
        self.hetero_encoder = to_hetero(
            self.encoder,
            metadata,
            aggr='mean'
        )

        # 分类器
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(hidden_channels, hidden_channels),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.5),
            torch.nn.Linear(hidden_channels, out_channels)
        )

    def forward(self, x_dict, edge_index_dict):
        # 通过异构图编码器
        x_dict = self.hetero_encoder(
            x_dict,
            edge_index_dict
        )

        # 获取用户节点表示
        user_x = x_dict['user']

        # 分类
        return self.classifier(user_x)


# ===================== 训练和评估函数 =====================
def train(model, loader, optimizer, criterion, device):
    """训练函数"""
    model.train()
    total_loss = 0
    total_correct = 0
    total_samples = 0

    for batch in tqdm(loader, desc="Training", leave=False):
        # 将batch移动到设备
        batch = batch.to(device)
        optimizer.zero_grad()

        # 前向传播
        out = model(
            batch.x_dict,
            batch.edge_index_dict
        )

        # 获取训练掩码
        train_mask = batch['user'].train_mask

        # 计算损失（仅训练节点）
        loss = criterion(out[train_mask], batch['user'].y[train_mask])

        # 反向传播
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

        # 计算准确率
        pred = out.argmax(dim=1)
        correct = (pred[train_mask] == batch['user'].y[train_mask]).sum().item()
        total_correct += correct
        total_samples += train_mask.sum().item()

    avg_loss = total_loss / len(loader)
    accuracy = total_correct / total_samples if total_samples > 0 else 0
    return avg_loss, accuracy


def evaluate(model, loader, mask_type, device):
    """评估函数"""
    model.eval()
    total_correct = 0
    total_samples = 0

    with torch.no_grad():
        for batch in tqdm(loader, desc=f"Evaluating ({mask_type})", leave=False):
            batch = batch.to(device)
            # 前向传播
            out = model(
                batch.x_dict,
                batch.edge_index_dict
            )
            # 选择掩码类型
            if mask_type == 'train':
                mask = batch['user'].train_mask
            elif mask_type == 'val':
                mask = batch['user'].val_mask
            elif mask_type == 'test':
                mask = batch['user'].test_mask
            else:
                mask = None

            # 计算准确率
            if mask is not None:
                pred = out.argmax(dim=1)
                correct = (pred[mask] == batch['user'].y[mask]).sum().item()
                total_correct += correct
                total_samples += mask.sum().item()
    accuracy = total_correct / total_samples if total_samples > 0 else 0.0
    return accuracy


# =====================  主程序  =====================
if __name__ == "__main__":
    # 配置参数
    # 定义一些定义数据会使用都的参数
    config = read_yaml(file_path="./config.yaml")
    hereto_data = construct_graph_data_hetero(df=pd.read_pickle("./dataset.pkl"))
    num_classes = config["data_config"]["num_classes"]
    HIDDEN_CHANNELS = 64
    BATCH_SIZE = 32
    NEIGHBOR_SIZES = [10, 5]  # 邻居采样大小 [第一跳, 第二跳]
    EPOCHS = 2
    LEARNING_RATE = 0.005
    DEVICE = config["model_config"]["device"]
    BEST_MODEL_SAVE_PATH = "./results/models/best_model.pt"

    print(f"Using device: {DEVICE}")

    # 生成数据
    generator = ClickDataGenerator()
    hetero_data = generator.build_hetero_data()

    # 打印数据信息
    print("\nHeteroData summary:")
    print(hetero_data)
    print("Metadata:", hetero_data.metadata())
    print(f"Number of users: {len(hetero_data['user'].x)}")
    print(f"Number of pages: {len(hetero_data['page'].x)}")
    print(f"Number of edges: {hetero_data['user', 'clicks', 'page'].edge_index.size(1)}")

    # 创建数据加载器
    print("\nCreating data loaders...")

    # 训练集加载器
    train_loader = NeighborLoader(
        hetero_data,
        num_neighbors=NEIGHBOR_SIZES,
        input_nodes=('user', hetero_data['user'].train_mask),
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=4 if os.cpu_count() > 4 else 0,
        # replace=True,
    )

    # 验证集加载器
    val_loader = NeighborLoader(
        hetero_data,
        num_neighbors=NEIGHBOR_SIZES,
        input_nodes=('user', hetero_data['user'].val_mask),
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=2 if os.cpu_count() > 2 else 0,
    )

    # 测试集加载器
    test_loader = NeighborLoader(
        hetero_data,
        num_neighbors=NEIGHBOR_SIZES,
        input_nodes=('user', hetero_data['user'].test_mask),
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=2 if os.cpu_count() > 2 else 0
    )

    # 初始化模型
    print("Initializing model...")
    model = HeteroGNN(
        metadata=hetero_data.metadata(),
        hidden_channels=HIDDEN_CHANNELS,
        out_channels=num_classes,
        heads=4
    ).to(DEVICE)

    # 打印模型结构
    print("\nModel architecture:")
    print(model)

    # 优化器和损失函数
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
    criterion = torch.nn.CrossEntropyLoss()

    # 训练循环
    print(f"\nStarting training for {EPOCHS} epochs...")
    best_val_acc = 0.0
    history = {'train_loss': [], 'train_acc': [], 'val_acc': []}

    for epoch in range(1, EPOCHS + 1):
        print(f"\nEpoch {epoch}/{EPOCHS}")
        # 训练阶段
        train_loss, train_acc = train(model, train_loader, optimizer, criterion, DEVICE)
        # 验证阶段
        val_acc = evaluate(model, val_loader, 'val', DEVICE)
        # 保存历史记录
        history['train_loss'].append(train_loss)
        history['train_acc'].append(train_acc)
        history['val_acc'].append(val_acc)
        # 保存最佳模型
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), BEST_MODEL_SAVE_PATH)
            print(f"  New best model saved with val acc: {val_acc:.4f}")
        # 打印训练统计信息
        print(f"  Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Acc: {val_acc:.4f}")

    # 加载最佳模型并测试
    model.load_state_dict(torch.load(BEST_MODEL_SAVE_PATH))
    test_acc = evaluate(model, test_loader, 'test', DEVICE)
    print(f"\nTest Accuracy: {test_acc:.4f}")
    print("\nTraining completed!")
